/*
 *                                  _              _
 *          _ __ _______ _ __ ___ | |_ ___   ___ | |
 *         | '_ \_  / _ \ '_ ` _ \| __/ _ \ / _ \| |
 *         | |_) / /  __/ | | | | | || (_) | (_) | |
 *         | .__/___\___|_| |_| |_|\__\___/ \___/|_|
 *         |_|
 *
 *         by hakanai
 */

#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <termios.h>
#include <stdint.h>
#include <time.h>
#include <stdbool.h>

#define PZEMTOOL_VER "v0.1.2"
#define DEFAULT_PORT "/dev/ttyUSB0"

#define REG_VOLTAGE     0x0000
#define REG_CURRENT_L   0x0001
#define REG_CURRENT_H   0X0002
#define REG_POWER_L     0x0003
#define REG_POWER_H     0x0004
#define REG_ENERGY_L    0x0005
#define REG_ENERGY_H    0x0006
#define REG_FREQUENCY   0x0007
#define REG_PF          0x0008
#define REG_ALARM       0x0009

#define PZEM_DEFAULT_ADDR   0xF8

#define CMD_RHR         0x03
#define CMD_RIR         0X04
#define CMD_WSR         0x06
#define CMD_CAL         0x41
#define CMD_REST        0x42

#define WREG_ALARM_THR   0x0001
#define WREG_ADDR        0x0002

#define UPDATE_TIME     200
#define READ_TIMEOUT 2

#define DEFAULT_DEVICE_ADDR (0xf8)

// {{{ crc16
static const uint16_t crcTable[] = {
  0X0000, 0XC0C1, 0XC181, 0X0140, 0XC301, 0X03C0, 0X0280, 0XC241,
  0XC601, 0X06C0, 0X0780, 0XC741, 0X0500, 0XC5C1, 0XC481, 0X0440,
  0XCC01, 0X0CC0, 0X0D80, 0XCD41, 0X0F00, 0XCFC1, 0XCE81, 0X0E40,
  0X0A00, 0XCAC1, 0XCB81, 0X0B40, 0XC901, 0X09C0, 0X0880, 0XC841,
  0XD801, 0X18C0, 0X1980, 0XD941, 0X1B00, 0XDBC1, 0XDA81, 0X1A40,
  0X1E00, 0XDEC1, 0XDF81, 0X1F40, 0XDD01, 0X1DC0, 0X1C80, 0XDC41,
  0X1400, 0XD4C1, 0XD581, 0X1540, 0XD701, 0X17C0, 0X1680, 0XD641,
  0XD201, 0X12C0, 0X1380, 0XD341, 0X1100, 0XD1C1, 0XD081, 0X1040,
  0XF001, 0X30C0, 0X3180, 0XF141, 0X3300, 0XF3C1, 0XF281, 0X3240,
  0X3600, 0XF6C1, 0XF781, 0X3740, 0XF501, 0X35C0, 0X3480, 0XF441,
  0X3C00, 0XFCC1, 0XFD81, 0X3D40, 0XFF01, 0X3FC0, 0X3E80, 0XFE41,
  0XFA01, 0X3AC0, 0X3B80, 0XFB41, 0X3900, 0XF9C1, 0XF881, 0X3840,
  0X2800, 0XE8C1, 0XE981, 0X2940, 0XEB01, 0X2BC0, 0X2A80, 0XEA41,
  0XEE01, 0X2EC0, 0X2F80, 0XEF41, 0X2D00, 0XEDC1, 0XEC81, 0X2C40,
  0XE401, 0X24C0, 0X2580, 0XE541, 0X2700, 0XE7C1, 0XE681, 0X2640,
  0X2200, 0XE2C1, 0XE381, 0X2340, 0XE101, 0X21C0, 0X2080, 0XE041,
  0XA001, 0X60C0, 0X6180, 0XA141, 0X6300, 0XA3C1, 0XA281, 0X6240,
  0X6600, 0XA6C1, 0XA781, 0X6740, 0XA501, 0X65C0, 0X6480, 0XA441,
  0X6C00, 0XACC1, 0XAD81, 0X6D40, 0XAF01, 0X6FC0, 0X6E80, 0XAE41,
  0XAA01, 0X6AC0, 0X6B80, 0XAB41, 0X6900, 0XA9C1, 0XA881, 0X6840,
  0X7800, 0XB8C1, 0XB981, 0X7940, 0XBB01, 0X7BC0, 0X7A80, 0XBA41,
  0XBE01, 0X7EC0, 0X7F80, 0XBF41, 0X7D00, 0XBDC1, 0XBC81, 0X7C40,
  0XB401, 0X74C0, 0X7580, 0XB541, 0X7700, 0XB7C1, 0XB681, 0X7640,
  0X7200, 0XB2C1, 0XB381, 0X7340, 0XB101, 0X71C0, 0X7080, 0XB041,
  0X5000, 0X90C1, 0X9181, 0X5140, 0X9301, 0X53C0, 0X5280, 0X9241,
  0X9601, 0X56C0, 0X5780, 0X9741, 0X5500, 0X95C1, 0X9481, 0X5440,
  0X9C01, 0X5CC0, 0X5D80, 0X9D41, 0X5F00, 0X9FC1, 0X9E81, 0X5E40,
  0X5A00, 0X9AC1, 0X9B81, 0X5B40, 0X9901, 0X59C0, 0X5880, 0X9841,
  0X8801, 0X48C0, 0X4980, 0X8941, 0X4B00, 0X8BC1, 0X8A81, 0X4A40,
  0X4E00, 0X8EC1, 0X8F81, 0X4F40, 0X8D01, 0X4DC0, 0X4C80, 0X8C41,
  0X4400, 0X84C1, 0X8581, 0X4540, 0X8701, 0X47C0, 0X4680, 0X8641,
  0X8201, 0X42C0, 0X4380, 0X8341, 0X4100, 0X81C1, 0X8081, 0X4040
};

/*
 * Calculate the CRC16-Modbus for a buffer
 * Based on https://www.modbustools.com/modbus_crc16.html
 */
static uint16_t crc16(const uint8_t *data, uint16_t len) {
  uint8_t nTemp; // CRC table index
  uint16_t crc = 0xFFFF; // Default value

  while (len--) {
    nTemp = *data++ ^ crc;
    crc >>= 8;
    crc ^= (uint16_t)crcTable[nTemp];
  }
  return crc;
}

bool check_crc(const uint8_t *buf, uint16_t len) {
  if(len <= 2) // Sanity check
    return false;

  uint16_t crc = crc16(buf, len - 2); // Compute CRC of data
  return ((uint16_t)buf[len-2]  | (uint16_t)buf[len-1] << 8) == crc;
}

void set_crc(uint8_t *buf, uint16_t len){
  if(len <= 2) // Sanity check
    return;

  uint16_t crc = crc16(buf, len - 2); // CRC of data

  // Write high and low byte to last two positions
  buf[len - 2] = crc & 0xFF; // Low byte first
  buf[len - 1] = (crc >> 8) & 0xFF; // High byte second
}
// }}}

// {{{ set_port_opt()
static void set_portopt(int fd) {
  struct termios options;

/*  tcgetattr(fd, &options);*/
  memset(&options, 0, sizeof(struct termios));

  cfsetispeed(&options, B9600);
  cfsetospeed(&options, B9600);

  /* Set 8 bit */
  options.c_cflag &= ~CSIZE;
  options.c_cflag |= CS8;

  /* Set 1 stop */
  options.c_cflag &= ~CSTOPB;

  /* Set no parity */
  options.c_cflag &= ~PARENB;
  options.c_iflag &= ~INPCK;

  /* Raw input */
  options.c_lflag &= ~(ICANON | ECHO | ECHOE | ISIG);

  /* Software flow control is disabled */
  options.c_iflag &= ~(IXON | IXOFF | IXANY);

  /* Raw ouput */
  options.c_oflag &=~ OPOST;

  /* Unused because we use open with the NDELAY option */
  options.c_cc[VMIN] = 0;
  options.c_cc[VTIME] = 0;

  if (tcsetattr(fd, TCSANOW, &options) < 0) {
    printf("Unable to set port options\n");
    exit(1);
  }
}
// }}}

static void printbuf(uint8_t* buffer, uint16_t len) {
  for(uint16_t i = 0; i < len; i++){
    printf("%.2x ", buffer[i]);
  }
}

uint16_t receive(int fd, uint8_t *resp, uint16_t len) {
  time_t startTime = time(NULL); // Start time for Timeout
  uint8_t c[1];
  uint8_t index = 0; // Bytes we have read

  while ((index < len) && ((time(NULL) - startTime) < READ_TIMEOUT)) {
    if (read(fd, c, 1) > 0) {
      resp[index++] = c[0];
    }
  }

  printf("recv: ");
  printbuf(resp, len);
  printf("\n");

/*  printf("received: %d\n", index);*/
  // Check CRC with the number of bytes read
  if(!check_crc(resp, index))
    return 0;

  return index;
}

static bool sendcmd(int fd, uint8_t cmd, uint16_t rAddr, uint16_t val, bool check, uint16_t slave_addr){
  uint8_t sendBuffer[8] = {} ; // Send buffer
  uint8_t respBuffer[8] = {} ; // Response buffer (only used when check is true)

  sendBuffer[0] = slave_addr;                   // Set slave address
  sendBuffer[1] = cmd;                     // Set command

  sendBuffer[2] = (rAddr >> 8) & 0xFF;     // Set high byte of register address
  sendBuffer[3] = (rAddr) & 0xFF;          // Set low byte =//=

  sendBuffer[4] = (val >> 8) & 0xFF;       // Set high byte of register value
  sendBuffer[5] = (val) & 0xFF;            // Set low byte =//=

  set_crc(sendBuffer, 8);                   // Set CRC of frame

  write(fd, sendBuffer, 8); // send frame

  printf("sent: ");
  printbuf(sendBuffer, 8);
  printf("\n");

  if (check && !receive(fd, respBuffer, 8)) { // if check enabled, read the response
    return false;
  }

  return true;
}

static uint8_t readAddress(int fd) {
  static uint8_t response[8];
  uint8_t addr = 0;

  // Read 1 register
  if (!sendcmd(fd, CMD_RHR, WREG_ADDR, 0x01, false, 0xf8))
    return 0;

  if(receive(fd, response, 8) < 1)  // Something went wrong
    return 0;

  // Get the current address
  addr = ((uint32_t)response[3] << 8 | // Raw address
      (uint32_t)response[4]);

  return addr;
}

static bool setAddress(int fd, uint8_t addr, uint8_t newaddr) {
  if (addr < 0x01 || addr > 0xF7) // sanity check
      return false;

  // Write the new address to the address register
  if(!sendcmd(fd, CMD_WSR, WREG_ADDR, newaddr, true, addr))
      return false;

  return true;
}

static void search(int fd) {
  static uint8_t response[7];

  for(uint16_t addr = 0x01; addr <= 0xF8; addr++){
    //Serial.println(addr);
    sendcmd(fd, CMD_RIR, 0x00, 0x01, false, addr);

    if(receive(fd, response, 7) != 7){ // Something went wrong
      continue;
    } else {
      printf("Found addr: %u\n", addr);
    }
  }
}

static int change_address(int fd, uint8_t addr) {
  uint8_t curaddr = 0x00;

  curaddr = readAddress(fd);
  printf("Device address: 0x%02x\n", curaddr);

  printf("Trying to change...\n");
  if (setAddress(fd, curaddr, addr))
    printf("Success.\n");
  else
    return false;

  addr = readAddress(fd);
  printf("Current device address is: 0x%02x\n", addr);

  return addr;
}

static int reset_energy(int fd, uint8_t addr) {
  uint8_t buffer[] = {0x00, CMD_REST, 0x00, 0x00};
  uint16_t length = 0;
  uint8_t reply[5];

  buffer[0] = addr;

  set_crc(buffer, 4);

  write(fd, buffer, 4);

  length = receive(fd, reply, 5);

  if(length == 0 || length == 5){
    printf("Reset energy for 0x%02x failed\n", addr);
    return false;
  }

  return true;

}

static void print_values(int fd, uint8_t addr) {
  static uint8_t response[25];
  float voltage, current, power, energy, frequency, pf;
  uint16_t alarms;

  sendcmd(fd, CMD_RIR, 0x00, 0x0A, false, addr);

  if(receive(fd, response, 25) != 25){ // Something went wrong
    printf("Error.\n");
    return;
  }

  // Update the current values
  voltage = ((uint32_t)response[3] << 8 | // Raw voltage in 0.1V
                              (uint32_t)response[4])/10.0;

  current = ((uint32_t)response[5] << 8 | // Raw current in 0.001A
                              (uint32_t)response[6] |
                              (uint32_t)response[7] << 24 |
                              (uint32_t)response[8] << 16) / 1000.0;

  power =   ((uint32_t)response[9] << 8 | // Raw power in 0.1W
                              (uint32_t)response[10] |
                              (uint32_t)response[11] << 24 |
                              (uint32_t)response[12] << 16) / 10.0;

  energy =  ((uint32_t)response[13] << 8 | // Raw Energy in 1Wh
                              (uint32_t)response[14] |
                              (uint32_t)response[15] << 24 |
                              (uint32_t)response[16] << 16) / 1000.0;

  frequency=((uint32_t)response[17] << 8 | // Raw Frequency in 0.1Hz
                              (uint32_t)response[18]) / 10.0;

  pf =      ((uint32_t)response[19] << 8 | // Raw pf in 0.01
                              (uint32_t)response[20])/100.0;

  alarms =  ((uint32_t)response[21] << 8 | // Raw alarm value
                              (uint32_t)response[22]);

  printf("\n");
  printf("Voltage:\t%fV\n", voltage);
  printf("Current:\t%fA\n", current);
  printf("Frequency:\t%fHz\n", frequency);
  printf("Power:\t\t%fW\n", power);
  printf("Energy:\t\t%fWh\n", energy);
  printf("pf:\t\t%f\n", pf);
  printf("alarms:\t\t%u\n", alarms);
  printf("\n");
}

static void usage(char *app) {
  printf("Usage: %s [options]\n\n", app);
  printf(" -t[addr]\t\tspecify target addr\n");
  printf(" -p[port]\t\ttty port (default /dev/ttyUSB0)\n");
  printf(" -s[num]\t\tchange device address.\n");
  printf(" -r[time]\t\tread values from device\n");
  printf(" -e\t\t\treset energy\n");
  printf("\n");
  printf("pzemtool %s by @hakanai\n", PZEMTOOL_VER);
  exit(1);
}

int main(int ac, char **av) {
  char *port = strdup(DEFAULT_PORT);
  int fd;
  int c;
  char *endptr;
  bool sflag = false;
  uint8_t saddr = 0x00;
  bool aflag = false;
  bool rflag = false;
  bool tflag = false;
  bool eflag = false;
  int rtime;

  while ((c = getopt(ac, av, "t:p:s:r:eahv")) != -1) {
    switch (c) {
      case 't':
        tflag = true;
        saddr = strtol(optarg, &endptr, 16);
        break;
      case 'p':
        free(port);
        port = strdup(optarg);
        break;
      case 'a':
        aflag = true;
        break;
      case 's':
        sflag = true;
        saddr = strtol(optarg, &endptr, 16);
        break;
      case 'r':
        rflag = true;
        rtime = atoi(optarg);
        if (rtime < 1)
          rtime = 1;
        break;
      case 'e':
        eflag = true;
        break;
      default:
        usage(av[0]);
    }
  }

  fd = open(port, O_RDWR | O_NOCTTY| O_EXCL | O_NONBLOCK);

  if (fd == -1) {
    perror(port);
    exit(2);
  }

  set_portopt(fd);

  if (aflag) {
    uint8_t curaddr = readAddress(fd);
    printf("Device address: 0x%02x\n", curaddr);
  }

  if (eflag) {
    uint8_t curaddr = saddr;

    if (!tflag) {
      printf("Requst device address\n");
      curaddr = readAddress(fd);
    }

    printf("Using device: 0x%02x\n", curaddr);

    if (reset_energy(fd, curaddr)) {
      printf("Energy reset successful\n");
    } else {
      printf("Error to reset energy for 0x%02x\n", curaddr);
    }
  }

  if (!tflag && sflag) {
    printf("Change device address to 0x%02x\n", saddr);
    change_address(fd, saddr);
  }

  if (rflag) {
    uint8_t curaddr = saddr;

    if (!tflag) {
      printf("Requst device address\n");
      curaddr = readAddress(fd);
    }

    printf("Using device: 0x%02x\n", curaddr);

    while (1) {
      print_values(fd, curaddr);
      sleep(rtime);
    }
  }

  close(fd);

  return 0;
}

